# --- START OF FILE multi_run_FINAL_FIX.py ---

import open3d as o3d
import numpy as np
import copy
import time
import glob
import os

# --- Configuration (Centralized) ---
class Config:
    # --- Paths ---
    PLY_DIRECTORY = "/home/xfzhou/workspace/data/exp1/s2_denoise" 
    PLY_PATTERN = os.path.join(PLY_DIRECTORY, "*.ply")
    OUTPUT_FILENAME = "final_model_pgo.ply"

    # --- Preprocessing & Segmentation ---
    SCALE_THRESHOLD = 10.0
    PLANE_DIST_THRESHOLD = 0.01

    # --- Pairwise Registration Parameters ---
    RANSAC_VOXEL_SIZE = 0.01
    ICP_VOXEL_SIZES = [0.005, 0.001, 0.0005]
    ICP_MAX_CORRESPONDENCE_DIST = [s * 1.5 for s in ICP_VOXEL_SIZES]

    # --- Final Model Post-processing ---
    FINAL_MODEL_VOXEL_SIZE = 0.0005

    # --- 配准质量控制 ---
    MIN_RANSAC_FITNESS = 0.6

# --- Helper & Interactive Functions ---

def load_ply_point_cloud(filepath: str) -> o3d.geometry.PointCloud:
    print(f"正在加载 PLY 文件: {filepath} ...")
    pcd = o3d.io.read_point_cloud(filepath)
    if not pcd.has_points():
        print(f"错误: 文件 {filepath} 为空或加载失败。")
        return None
    print(f"加载完成，包含 {len(pcd.points)} 个点。")
    return pcd

def check_and_scale_to_meters(pcd: o3d.geometry.PointCloud, scale_threshold=10.0):
    max_bound = pcd.get_max_bound()
    if np.any(max_bound > scale_threshold):
        print(f"检测到点云单位可能为毫米（最大坐标: {max_bound.max():.2f} > {scale_threshold}）。")
        print("--> 正在将点云从毫米缩放到米（除以1000）...")
        pcd.scale(0.001, center=(0, 0, 0))
    return pcd

def _find_plane_from_picked_points(pcd: o3d.geometry.PointCloud, 
                                   indices: list, 
                                   dist_threshold: float) -> (np.ndarray, list):
    if len(indices) < 3: return None, None
    picked_points = pcd.select_by_index(indices).points
    centroid = np.mean(picked_points, axis=0)
    points_centered = picked_points - centroid
    _, _, vh = np.linalg.svd(points_centered)
    normal = vh[2, :]
    d = -np.dot(normal, centroid)
    plane_model = np.array([normal[0], normal[1], normal[2], d])
    all_points = np.asarray(pcd.points)
    distances = np.abs(np.dot(all_points, plane_model[:3]) + plane_model[3])
    inliers = np.where(distances < dist_threshold)[0].tolist()
    return plane_model, inliers

def select_and_remove_planes_interactive(pcd: o3d.geometry.PointCloud, 
                                         pcd_name: str, 
                                         plane_dist_threshold: float) -> o3d.geometry.PointCloud:
    pcd_to_edit = copy.deepcopy(pcd)
    while True:
        print("\n" + "="*50)
        print(f"交互式【平面】移除: {pcd_name}")
        print("请在弹出的窗口中选择要移除的平面上的至少3个点:")
        print("  - [Shift + 鼠标左键]: 旋转/平移")
        print("  - [Ctrl + 鼠标左键]: 点选点 (点会变成粉红色)")
        print("  - 选好点后，请按 'q' 关闭窗口继续。")
        vis = o3d.visualization.VisualizerWithEditing()
        vis.create_window(window_name=f"为 {pcd_name} 选择平面点 (选好后按 q)")
        vis.add_geometry(pcd_to_edit)
        vis.run()
        vis.destroy_window()
        picked_indices = vis.get_picked_points()
        if len(picked_indices) < 3:
            print("选择的点少于3个。是否结束平面移除？(y/n)")
            if input("> ").lower() == 'y': break
            else: continue
        _, inliers = _find_plane_from_picked_points(pcd_to_edit, picked_indices, plane_dist_threshold)
        if inliers and len(inliers) > 0:
            print(f"成功识别平面！包含 {len(inliers)} 个点。")
            pcd_plane = pcd_to_edit.select_by_index(inliers)
            pcd_object = pcd_to_edit.select_by_index(inliers, invert=True)
            pcd_plane.paint_uniform_color([0.8, 0.1, 0.1])
            pcd_object.paint_uniform_color([0.1, 0.8, 0.1])
            print("正在显示移除预览... 请确认是否移除该平面。(y/n)")
            o3d.visualization.draw_geometries([pcd_plane, pcd_object], window_name="平面移除预览")
            if input("> ").lower() == 'y':
                pcd_to_edit = pcd_object
                pcd_to_edit.paint_uniform_color([0.5, 0.5, 0.5])
                print("平面已移除。")
            else:
                print("操作已取消。")
        else:
            print("无法根据您的选择清晰地识别出一个平面。请重试。")
        print("是否要继续移除其他平面？(y/n)")
        if input("> ").lower() != 'y': break
    print(f"交互式平面移除完成。剩余点数: {len(pcd_to_edit.points)}")
    print("="*50 + "\n")
    return pcd_to_edit

# ==============================================================================
# === 最终修复版功能: 交互式移除套索（Lasso）选择的点 ===
# ==============================================================================
def select_and_remove_lasso_points_interactive(pcd: o3d.geometry.PointCloud, 
                                               pcd_name: str) -> o3d.geometry.PointCloud:
    """
    通过Open3D的顶点选择功能，交互式地移除一个或多个套索区域内的点。
    此版本使用稳定可靠的 draw_geometries_with_vertex_selection 方法。
    """
    pcd_to_edit = copy.deepcopy(pcd)
    
    while True:
        print("\n" + "="*50)
        print(f"交互式【套索】移除: {pcd_name}")
        print("一个窗口将弹出，请按以下步骤操作以选择【要移除】的点：")
        print("  1. 调整好视角。")
        print("  2. 按住 [Shift] 并用 [鼠标左键] 拖拽，画一个多边形来框选点。")
        print("     (被选中的点会变为黄色)")
        print("  3. 完成选择后，按下 [Q] 键关闭窗口。")
        print("  (如果不想移除任何点，直接按 [Q] 关闭窗口即可)")

        # 这个函数会打开一个窗口，并等待用户选择。
        # 当用户关闭窗口后，它会返回一个包含所有被选中点索引的列表。
        vis = o3d.visualization.VisualizerWithVertexSelection()
        vis.create_window(window_name=f"为 {pcd_name} 套索选择要移除的点 (Shift+拖拽 -> Q)")
        vis.add_geometry(pcd_to_edit)
        vis.run()
        picked_indices_dict = vis.get_picked_points()
        vis.destroy_window()

        if not picked_indices_dict:
            print("未选择任何点。是否结束套索移除？(y/n)")
            if input("> ").lower() == 'y':
                break
            else:
                continue
        
        # get_picked_points() 返回的是一个字典列表，我们需要提取索引
        picked_indices = [p.index for p in picked_indices_dict]

        # 分离点云用于预览
        pcd_object = pcd_to_edit.select_by_index(picked_indices, invert=True)
        pcd_removed = pcd_to_edit.select_by_index(picked_indices)

        print(f"成功用套索选中 {len(pcd_removed.points)} 个点用于移除。")
        pcd_object.paint_uniform_color([0.1, 0.8, 0.1])  # 绿色: 保留的点
        pcd_removed.paint_uniform_color([0.8, 0.1, 0.1]) # 红色: 被移除的点
        
        print("正在显示移除预览... 请确认是否移除该区域的点。(y/n)")
        o3d.visualization.draw_geometries([pcd_removed, pcd_object], window_name="套索移除预览")
        
        if input("> ").lower() == 'y':
            pcd_to_edit = pcd_object
            pcd_to_edit.paint_uniform_color([0.5, 0.5, 0.5]) # 恢复默认颜色
            print("套索区域的点已移除。")
        else:
            print("操作已取消。")

        print("是否要继续移除其他区域？(y/n)")
        if input("> ").lower() != 'y':
            break
            
    print(f"交互式套索移除完成。点数: {len(pcd.points)} -> {len(pcd_to_edit.points)}")
    print("="*50 + "\n")
    return pcd_to_edit

def draw_transformed_points(points, transformations):
    points_list = []
    for p,t in zip(points, transformations):
        p_tmp = copy.deepcopy(p)
        if t is not None: p_tmp.transform(t)
        points_list.append(p_tmp)
    o3d.visualization.draw_geometries(points_list, window_name="结果展示")
    return points_list

def run_pairwise_registration(source_obj, target_obj, config):
    print("--- 正在执行成对配准 (带质量检查) ---")
    source_down, source_fpfh = preprocess_for_ransac(source_obj, config.RANSAC_VOXEL_SIZE)
    target_down, target_fpfh = preprocess_for_ransac(target_obj, config.RANSAC_VOXEL_SIZE)
    distance_threshold = config.RANSAC_VOXEL_SIZE * 1.5
    ransac_result = o3d.pipelines.registration.registration_ransac_based_on_feature_matching(
        source_down, target_down, source_fpfh, target_fpfh, True, distance_threshold,
        o3d.pipelines.registration.TransformationEstimationPointToPoint(False), 3, 
        [o3d.pipelines.registration.CorrespondenceCheckerBasedOnEdgeLength(0.9),
         o3d.pipelines.registration.CorrespondenceCheckerBasedOnDistance(distance_threshold)], 
        o3d.pipelines.registration.RANSACConvergenceCriteria(100000, 0.999))
    if ransac_result.fitness < config.MIN_RANSAC_FITNESS:
        print(f"[警告] RANSAC 配准失败！Fitness ({ransac_result.fitness:.3f}) < 阈值 ({config.MIN_RANSAC_FITNESS:.3f}).")
        return None, None
    print(f"RANSAC 成功, Fitness: {ransac_result.fitness:.3f}")
    icp_result = run_fine_icp_on_full_cloud(
        source_obj, target_obj, ransac_result.transformation,
        config.ICP_VOXEL_SIZES, config.ICP_MAX_CORRESPONDENCE_DIST)
    information_matrix = o3d.pipelines.registration.get_information_matrix_from_point_clouds(
        source_obj, target_obj, config.ICP_MAX_CORRESPONDENCE_DIST[-1], icp_result.transformation)
    return icp_result.transformation, information_matrix

def preprocess_for_ransac(pcd, voxel_size):
    pcd_down = pcd.voxel_down_sample(voxel_size)
    radius_normal = voxel_size * 2
    pcd_down.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=radius_normal, max_nn=30))
    radius_feature = voxel_size * 5
    pcd_fpfh = o3d.pipelines.registration.compute_fpfh_feature(
        pcd_down, o3d.geometry.KDTreeSearchParamHybrid(radius=radius_feature, max_nn=100))
    return pcd_down, pcd_fpfh

def run_fine_icp_on_full_cloud(source_full, target_full, initial_transform, voxel_sizes, max_corr_dists):
    current_transform = initial_transform
    for i, voxel_size in enumerate(voxel_sizes):
        max_corr_dist = max_corr_dists[i]
        source_down = source_full.voxel_down_sample(voxel_size)
        target_down = target_full.voxel_down_sample(voxel_size)
        source_down.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=voxel_size * 2, max_nn=30))
        target_down.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=voxel_size * 2, max_nn=30))
        reg_p2l = o3d.pipelines.registration.registration_icp(
            source_down, target_down, max_corr_dist, current_transform,
            o3d.pipelines.registration.TransformationEstimationPointToPlane())
        current_transform = reg_p2l.transformation
    return reg_p2l

# --- Main Orchestrator ---

def main():
    cfg = Config()
    start_time_total = time.time()

    pcd_files = glob.glob(cfg.PLY_PATTERN)
    pcd_files = sorted(pcd_files, key=lambda x:int(os.path.basename(x)[1:-4]))
    if not pcd_files:
        print(f"错误: 在目录 '{cfg.PLY_DIRECTORY}' 中没有找到匹配 '{cfg.PLY_PATTERN}' 的文件。")
        return
        
    print(f"找到 {len(pcd_files)} 个点云文件。")
    
    pcd_obj_list = []
    for i, fpath in enumerate(pcd_files):
        print(f"\n--- 正在处理文件 {i+1}/{len(pcd_files)}: {os.path.basename(fpath)} ---")
        pcd_full = load_ply_point_cloud(fpath)
        if pcd_full is None: continue
        pcd_obj = check_and_scale_to_meters(pcd_full, cfg.SCALE_THRESHOLD)
        pcd_obj_list.append(pcd_obj)

    # --- Phase 2: 构建位姿图 ---
    print("\n" + "#"*60)
    print("--- Phase 2: 正在构建位姿图 (Pose Graph) ---")
    print("#"*60)
    
    n_pcds = len(pcd_obj_list)
    transformation_icp_list = []
    for i in range(n_pcds - 1):
        source_idx, target_idx = i + 1, i
        print(f"\n... 正在配准 Scan {source_idx} -> Scan {target_idx} ...")
        transformation_icp, _ = run_pairwise_registration(pcd_obj_list[source_idx], pcd_obj_list[target_idx], cfg)
        if transformation_icp is None:
            print(f"\n\n{'!'*20} 致命错误 {'!'*20}")
            print(f"无法将 Scan {source_idx} 与 Scan {target_idx} 对齐。处理中止。")
            return
        transformation_icp_list.append(transformation_icp)

    transformations = [np.identity(4)]
    for i in range(n_pcds - 1):
        new_transform = np.dot(transformations[-1], transformation_icp_list[i])
        transformations.append(new_transform)

    pcd_combined = o3d.geometry.PointCloud()
    for i in range(n_pcds):
        pcd_obj_list[i].transform(transformations[i])
        pcd_combined += pcd_obj_list[i]
        
    print("\n正在对组合后的模型进行体素降采样以获得均匀密度...")
    pcd_combined_down = pcd_combined.voxel_down_sample(cfg.FINAL_MODEL_VOXEL_SIZE)
    
    o3d.io.write_point_cloud(cfg.OUTPUT_FILENAME, pcd_combined_down, write_ascii=False)
    total_time = time.time() - start_time_total
    print("\n" + "="*60)
    print(f"所有计算完成，总耗时: {total_time:.2f} 秒。")
    print(f"最终的完整模型已保存至: {cfg.OUTPUT_FILENAME}")
    print("="*60)
    
    print("\n正在显示最终的组合模型...")
    o3d.visualization.draw_geometries([pcd_combined_down], window_name="Final Registered Model")

if __name__ == "__main__":
    main()

# --- END OF FILE ---